from langchain_core.runnables.utils import ConfigurableField
from langchain_openai import ChatOpenAI
# from langchain_deepseek import ChatDeepSeek
from langchain_community.chat_models import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.schema import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
import os

# 模型-openai
gpt_3_5_turbo_model = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

# 模型-deepseek
gpt_4o_mini_model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# 通过 configurable_alternatives 按指定字段选择模型
model = gpt_4o_mini_model.configurable_alternatives(
    ConfigurableField(id="llm"), 
    default_key="gpt_4o_mini", 
    gpt_3_5=gpt_3_5_turbo_model,
    # claude=claude_model,
)

# Prompt 模板
prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template("{query}"),
    ]
)

# LCEL
chain = (
    {"query": RunnablePassthrough()} 
    | prompt
    | model 
    | StrOutputParser()
)

# 运行时指定模型 "gpt" or "ernie"
ret = chain.with_config(configurable={"llm": "gpt_3_5"}).invoke("请自我介绍")
print(ret.content)